from major.models_manager import embedding_model, chat_model
from langchain_chroma import Chroma

# 需求：基于上一步构建好的文档库，构建一个rag系统，实现基于 rag 的问答

# 加载文档库
vector_db = Chroma(persist_directory = './chroma_db',
       embedding_function = embedding_model.get_model())

llm = chat_model.get_model()
# 检索与生成
def generate(query):
    # 检索一下与query相关的文档
    docs = vector_db.similarity_search( query, k=3)
    docs = [f"简历 {i+1}:\n\n[[{doc.page_content}]]" for i,doc in enumerate(docs)]
    doc_context_str = "\n\n".join(docs)

    # 拼接prompt
    prompt = f'''
    请根据相关简历推荐候选人
    相关简历：
    {doc_context_str}
    
    用户问题：{query}
    
    '''
    res = llm.stream(prompt)
    for i in res:
        print(i.content, end='')

generate("请推荐一个做过CV的候选人")



